#include "header.h"

static fftwq_plan plan_br1, plan_br2;
static __float128 *k, *u, overN, q, *mk, dtau, gama, sp, num, den, err, err2, l2, tol, hmax, ml0, ml, htt;
static __float128 coeff, aa, bb, NV, *ra[4], *rt1, *rt2, *rhu[1], *rhk[2], *redk[2], *reDk, *reD;
static __float128 ip, *reDk, *reRk;
static __float128 cg_error;
static int N, n2, nthreads; 
static work_ptr tmp;
static pthread_t thread1, thread2, thread3;
static pthread_mutex_t w_mut;
static pthread_cond_t w_cond;
static int w_threads, nh, nq;

void err_msg(char* str)
{
  printf("%s\n",str);
  exit(1);
}


void init(int *argc, char **argv, params_ptr in, work_ptr wrk)
{
  char str[480], line[480], value[480], param[480];
  sprintf(str, "Usage:\n\t%s input", argv[0]);
  if (*argc != 2) err_msg(str);

  sprintf(str, "%s.cfg", argv[1]);
  FILE *fh = fopen(str,"r"); 
  FILE *fhlog = fopen("run.log","w"); 
  if (fh == NULL) err_msg("Cannot open file");
  while (fgets(line, 480, fh)!=NULL) {
    sscanf(line, "%s\t%s", param, value);
    if (strcmp(param,"#runname=") == 0) sprintf(in->runname,"%s", value);
    if (strcmp(param,"#resname=") == 0) sprintf(in->resname,"%s", value);
    if (strcmp(param,"#npoints=") == 0) in->n = atoi(value);
    if (strcmp(param,"#gravity=") == 0) in->g = strtoflt128 (value, NULL);
    if (strcmp(param,"#velocit=") == 0) in->c = strtoflt128 (value, NULL);
    if (strcmp(param,"#delta_t=") == 0) in->dtau = strtoflt128 (value, NULL);
    if (strcmp(param,"#toleran=") == 0) in->tol = strtoflt128 (value, NULL);
    if (strcmp(param,"#perturb=") == 0) in->alfa = strtoflt128 (value, NULL);
    if (strcmp(param,"#io_time=") == 0) in->t_skip = atoi(value);
    if (strcmp(param,"#nthread=") == 0) in->nthreads = atoi(value);
    if (strcmp(param,"#hmax_tt=") == 0) htt = strtoflt128 (value, NULL);
  }
  fprintf(fhlog, "runname = %s\n", in->runname );
  fprintf(fhlog, "resname = %s\n", in->resname );
  fprintf(fhlog, "npoints = %d\n", in->n );
  fprintf(fhlog, "gravity = %Qe\n", in->g );
  fprintf(fhlog, "velocit = %Qe\n", in->c );
  fprintf(fhlog, "toleran = %Qe\n", in->tol );
  fprintf(fhlog, "delta_t = %Qe\n", in->dtau );
  fprintf(fhlog, "perturb = %Qe\n", in->alfa );
  fprintf(fhlog, "hmax_tt = %Qe\n", htt );
  fprintf(fhlog, "io_time = %d\n", in->t_skip );
  fprintf(fhlog, "nthread = %d\n", in->nthreads );
  fclose(fh); fclose(fhlog);
  N = in->n; n2 = N/2 + 1; nthreads = in->nthreads;

  if (nthreads != 1) {
   fftwq_init_threads();
   fftwq_plan_with_nthreads(nthreads);
  }

  ra[0] = fftwq_malloc(sizeof(__float128)*n2);
  ra[1] = fftwq_malloc(sizeof(__float128)*n2);
  ra[2] = fftwq_malloc(sizeof(__float128)*n2);
  ra[3] = fftwq_malloc(sizeof(__float128)*n2);
  wrk->YY  = fftwq_malloc(sizeof(__float128)*n2);
  wrk->YYk = fftwq_malloc(sizeof(__float128)*n2);

  rhu[0] = fftwq_malloc(sizeof(__float128)*n2);
  rhk[0] = fftwq_malloc(sizeof(__float128)*n2);
  rhk[1] = fftwq_malloc(sizeof(__float128)*n2);

  redk[0] = fftwq_malloc(sizeof(__float128)*n2);
  redk[1] = fftwq_malloc(sizeof(__float128)*n2);

  reDk = fftwq_malloc(sizeof(__float128)*n2);
  reD  = fftwq_malloc(sizeof(__float128)*n2);
  reRk  = fftwq_malloc(sizeof(__float128)*n2);

  //u  = fftwq_malloc(sizeof(__float128)*(in->n));
  q  = (in->g)/powq(in->c,2);
  
  sprintf(value,"r2r.fftw.%d", nthreads);
  if (fftwq_import_wisdom_from_filename(value) == 0) {
    printf("Creating plans\n");
    fftwq_set_timelimit(10.);
    plan_br1 = fftwq_plan_r2r_1d(n2, ra[1], ra[0], FFTW_REDFT00, fmode);
    plan_br2 = fftwq_plan_r2r_1d(n2, ra[3], ra[2], FFTW_REDFT00, fmode);
    fftwq_export_wisdom_to_filename(value);
    err_msg("Plans created. Rerun the computation");
  } else {
    fftwq_set_timelimit(10.);
    plan_br1 = fftwq_plan_r2r_1d(n2, ra[1], ra[0], FFTW_REDFT00, fmode);
    plan_br2 = fftwq_plan_r2r_1d(n2, ra[3], ra[2], FFTW_REDFT00, fmode);
  }
  /*for (int j = 0; j < in->n; j ++) {
    u[j] = 2*pi*j/(in->n);
    wrk->YYk[j] = 0.1*exp(-2*pow(u[j]-pi,2)); // -cos(3.*u[j])
  }*/
  fh = fopen(in->resname,"r");
  int jj = 0; 
  //double nk = 0.;
  if (fh != NULL) {
    char v1[80],v2[80],v3[80],v4[80], *dummy;
    while (fgets(line, 480, fh)!=NULL) {
      if (jj == 0) {
        dummy = fgets(line, 160, fh);
        dummy = fgets(line, 160, fh);
        dummy = fgets(line, 160, fh);
      }
      sscanf(line, "%s\t%s\t%s\t%s\n",  v1, v2, v3, v4);
      //nk = atof(v2); 
      if (atof(v2) < -0.1) break;
      if (jj == n2) {
        break;
      }
      wrk->YYk[jj] = strtoflt128(v4, NULL);
      jj++;      
    }
    fclose(fh);
    printf("Lines read = %d of %d expected\n", jj, n2);
    if ( jj != n2) {
      printf("Restarting from coarser grid.\nMissing Fourier modes are set to zero.\n");
      for (int j = jj; j < n2; j++) wrk->YYk[j] = 0.Q;
    }
  }

  overN = 1.Q/N; 
  for(int j = 0; j < n2; j++) wrk->YYk[j] = (wrk->YYk[j])*cexpq((in->alfa)*j);
  in->w_ptr = wrk; dtau = in->dtau;
  gama = 1.Q + 1.Q/dtau; tol = 1.8e-32Q*sqrtq(N);
  ifft_even(wrk->YYk,wrk->YY);
  //for(int j = 0; j < n2; j++) wrk->YYk[j] = (wrk->YYk[j])*overN*cexpq((in->alfa)*j); 
  /*if (jj == 0) {
    wrk->Yk[0] = 0;
    for (int j = 1; j < N; j++) wrk->Yk[0] += -cabsq(k[j])*conj(wrk->Yk[j])*(wrk->Yk[j]); 
  }*/
}


void ifft_even(__float128 *in, __float128 *out) {
  memcpy(ra[1], in, n2*sizeof(__float128));
  fftwq_execute(plan_br1);
  memcpy(out, ra[0], n2*sizeof(__float128));
}

void ifft_even_two(__float128 *in, __float128 *out) {
  memcpy(ra[3], in, n2*sizeof(__float128));
  fftwq_execute(plan_br2);
  memcpy(out, ra[2], n2*sizeof(__float128));
}

__float128 inner_product_real(__float128 *in1, __float128 *in2) {
  ip = 0.Q;
  for (int j = 0; j < n2; j++) {
    ip += 2.Q*in1[j]*in2[j];
  }
  ip = ip - (in1[0]*in2[0] + in1[n2-1]*in2[n2-1]);
  return ip;
}

__float128 mi_product_real(__float128 *in1, __float128 *in2) {
  ip = 0.Q;
  for (int j = 0; j < n2; j++) {
    ip += 2.Q*in1[j]*in2[j]/(j - q);
  }
  ip = ip - (in1[0]*in2[0]/(-q) + in1[n2-1]*in2[n2-1]/(n2-1-q));
  return ip;
}

void prepare_rhs_even(work_ptr wrk) {
  ifft_even(wrk->YYk,wrk->YY); 
  for (int j = 0; j < n2; j++) {
    rhk[0][j] = (wrk->YYk[j])*j;
    rhk[1][j] = (wrk->YY[j])*(wrk->YY[j]);
  }
  ifft_even(rhk[1],rhk[1]);  
  ifft_even(rhk[0],rhu[0]);
  for (int j = 0; j < n2; j++) {
    rhk[1][j] = j*rhk[1][j]*overN;   
    rhk[0][j] = rhu[0][j]*(wrk->YY[j])*overN;
  }
  ifft_even(rhk[0],rhk[0]); 
  for (int j = 0; j < n2; j++) {
    rhk[0][j] = -(j - q)*(wrk->YYk[j]) + q*(0.5Q*rhk[1][j] + rhk[0][j]); 
  }
}

/*void *run_thread_one(void *in) {
  work_ptr WRK  = (work_ptr) in;
  //pthread_cond_wait(&w_cond, &w_mut);
  //pthread_mutex_unlock(&w_mut);
  pthread_mutex_lock(&w_mut);
  --w_threads;
  //pthread_cond_signal(&w_cond);  
  if (w_threads == 0) pthread_cond_broadcast(&w_cond);
  else pthread_cond_signal(&w_cond);
  pthread_cond_wait(&w_cond, &w_mut);  
  pthread_mutex_unlock(&w_mut); 
  printf("Complete\n");

}*/

void barrier_two(pthread_mutex_t *mut, pthread_cond_t *cond) {
  pthread_mutex_lock(mut);
  --w_threads;
  if (w_threads == 0) pthread_cond_broadcast(cond);
  else pthread_cond_signal(cond);
  pthread_cond_wait(cond, mut);
  pthread_mutex_unlock(mut);
}

void barrier_main(pthread_mutex_t *mut, pthread_cond_t *cond) {
  pthread_mutex_lock(mut);
  --w_threads;
  if (w_threads == 0) {
    pthread_cond_signal(cond);
  } else {
    pthread_cond_wait(cond, mut);
  }
  pthread_mutex_unlock(mut);
  w_threads = 2; 
  pthread_cond_signal(cond);
}

void *run_thread_two(void *in) {
  work_ptr WRK  = (work_ptr) in;
  //pthread_cond_wait(&w_cond, &w_mut);
  //pthread_mutex_unlock(&w_mut);

  for (int j = nh; j < n2; j++)  redk[1][j] = j*reDk[j]; 
  barrier_two(&w_mut, &w_cond);  //stop 1
  ifft_even_two(redk[1], redk[1]);
  barrier_two(&w_mut, &w_cond);  // stop 2
  for (int j = 0; j < n2; j++) redk[1][j] = overN*((WRK->YY[j])*redk[1][j] + rhu[0][j]*reD[j]);
  ifft_even_two(redk[1], redk[1]);
  barrier_two(&w_mut, &w_cond);  // stop 3 
  for (int j = nh; j < n2; j++) redk[0][j] = -(j - q)*reDk[j] + q*(j*overN*redk[0][j] + redk[1][j]);
  barrier_two(&w_mut, &w_cond);  // stop 3.5 
  den = inner_product_real(reDk,redk[0]);
  barrier_two(&w_mut, &w_cond);   // stop 4
  for (int j = 0; j < n2; j++) rhk[1][j] +=  aa*reDk[j]; 
  barrier_two(&w_mut, &w_cond);  // stop 5
  for (int j = nh; j < n2; j++) reRk[j]   += -aa*redk[0][j];
  barrier_two(&w_mut, &w_cond);  // stop 6
  cg_error = sqrtq(inner_product_real(reRk,reRk));
  barrier_two(&w_mut, &w_cond);  // stop 7
  pthread_exit(NULL);
}
void prepare_linearization_even(work_ptr wrk) {
  int iret1, iret2;
  void *aux = (void *)wrk;
  nh = (n2-1)/2;
  pthread_mutex_init(&w_mut,NULL);
  pthread_cond_init(&w_cond,NULL);
  w_threads = 2;  
  iret2 = pthread_create( &thread2, NULL, &run_thread_two, aux);
  for (int j = 0; j < nh; j++)  redk[1][j] = j*reDk[j]; 
  barrier_main(&w_mut, &w_cond);   // stop 1
  ifft_even(reDk, reD);  
  barrier_main(&w_mut, &w_cond);   // stop 2
  for (int j = 0; j < n2; j++) redk[0][j] = (wrk->YY[j])*reD[j]; 
  ifft_even(redk[0], redk[0]);
  barrier_main(&w_mut, &w_cond);   // stop 3
  for (int j = 0; j < nh; j++) redk[0][j] = -(j - q)*reDk[j] + q*(j*overN*redk[0][j] + redk[1][j]);
  barrier_main(&w_mut, &w_cond);   // stop 3.5
  num = mi_product_real(reRk,reRk);  
  barrier_main(&w_mut, &w_cond);   // stop 4
  aa = num/den;
  den = mi_product_real(reRk, reRk);  
  barrier_main(&w_mut, &w_cond);   // stop 5
  for (int j = 0; j < nh; j++) reRk[j]   += -aa*redk[0][j];
  num = -aa*mi_product_real(redk[0], reRk);  // Polak-Ribiere
  bb = num/den;
  barrier_main(&w_mut, &w_cond);   // stop 6
  for (int j = 0; j < n2; j++) reDk[j] = reRk[j]/(j - q) + bb*reDk[j];
  barrier_main(&w_mut, &w_cond);   // stop 7
  pthread_join(thread2, NULL);





  /*for (int j = 0; j < n2; j++) {
    redk[1][j] = j*reDk[j]; 
  }   
  ifft_even(reDk, reD);
  ifft_even(redk[1], redk[1]);

  for (int j = 0; j < n2; j++) {
    redk[0][j] = (wrk->YY[j])*reD[j];      
    redk[1][j] = overN*((wrk->YY[j])*redk[1][j] + rhu[0][j]*reD[j]); 
  }
  fft_even(redk[0], redk[0]);
  fft_even(redk[1], redk[1]);*/
}

void iterate_cg(work_ptr wrk) {  
  int counter, counter2;
  __float128 err = 1.Q;
  __float128 eps = 0.01Q;

  counter2 = 1; NV = 1.Q; counter = 0; cg_error = 1.Q;
  while (NV > tol) {
    prepare_rhs_even(wrk);
    writer_out(wrk->YYk,"restart.dat", NV, wrk->YY[0]);  
    //writer_out(wrk->YYk,"stokes_spec.dat");
    for (int j = 0; j < n2; j++) {
      rhk[1][j] = 0.Q;
      reRk[j] = -1.Q*rhk[0][j];
      reDk[j] = reRk[j]/(j - q);
    }
    NV = sqrtq(inner_product_real(rhk[0],rhk[0]) );  
    if ((counter2%1) == 0 ) {
        FILE *fhlog = fopen("run.log","a"); 
        hmax = ((wrk->YY[n2-1])-(wrk->YY[0]))/(2.Q*pi);
	printf("NV = %Qe Newton cycle %1d\n", NV, counter2-1);
	fprintf(fhlog, "NV = %Qe Newton cycle %1d\n", NV, counter2-1);
        fclose(fhlog);
    }
    counter = 0; cg_error = 1.Q;
    if (NV < 1e-28) eps = 0.25;
    while (cg_error > eps*NV ) {       // start CG cycle 
      prepare_linearization_even(wrk);
      counter++;
      if ((counter%5)==0) {
        FILE *fhlog = fopen("run.log","a"); 
        printf("CG Error = %Qe at %d step\n", cg_error, counter);
        fprintf(fhlog, "CG Error = %Qe at %d step\n", cg_error, counter);
        fclose(fhlog);
      }
      if (counter == 4000) break;

    }                                  // end of CG cycle
    if (counter == 4000) break;
    for (int j = 0; j < n2; j++) wrk->YYk[j] = wrk->YYk[j] + rhk[1][j];
    hmax = ((wrk->YY[n2-1])-(wrk->YY[0]))/(2.Q*pi);
    prepare_rhs_even(wrk);
    NV = sqrtq(inner_product_real(rhk[0],rhk[0]) );  
    printf("NV = %Qe at %3d CG step, H/L = %.32Qe Speed = %.18Qe\n", NV, counter, hmax, 1.Q/sqrtq(q));
    FILE *fhlog = fopen("run.log","a"); 
    fprintf(fhlog, "NV = %Qe at %3d CG step, H/L = %.32Qe Speed = %.18Qe\n", NV, counter, hmax, 1.Q/sqrtq(q));
    fclose(fhlog);
    counter2++;
   
  }  // end of outer cycle

  if (counter == 4000) err_msg("CG diverged");

  ifft_even(wrk->YYk,wrk->YY);  
  writer_out(wrk->YYk,"restart.dat", NV, wrk->YY[0]);
  FILE *fid = fopen("../steep.data","a");
  fprintf(fid, "%.32Qe\t%.32Qe\t%.32Qe\n", hmax, 1.Q/sqrtq(q), wrk->YY[0]);
  fclose(fid);
}

















